/*
 * SeparatedCGOptimizer.cpp
 *
 *      Author: Jerome RD Soine and Christoph A Brand
 *      Institution: Schwarz goup, Institute for Theoretical Physics, Heidelberg University, Germany
 */

#include "SeparatedCGOptimizer.h"
#include "include/core/utility/VectorOperations.h"
#include <iostream>

using namespace std;

void SeparatedCGOptimizer::run() {
	cout << "\n \n STARTING SEPARATED CG OPTIMIZATION \n \n";
	write_current_state(0);
	for(tot_iterations=1;tot_iterations<=n_iterations;tot_iterations++)
	{
		if( (first && tot_iterations%cg_block_size == 1) || (!first && (tot_iterations%cg_block_size) == cg_block_size) ){
			adjust_first_dof();
			current_CG_iteration_step = 0;
		} else {
			iterate_cg();
		}
		write_current_state(tot_iterations);
	}
}

void SeparatedCGOptimizer::adjust_first_dof(){

	cout << " START NETWORK OPTIMIZATION STEP" << endl << endl;

	vector<double> displaced_pos = solution;
	displaced_pos[0] += gradient_spacing_base;
	Worker * reference = worker_manager->getWorker(solution, -1);
	Worker * displaced = worker_manager->getWorker(displaced_pos, -1);
	reference->start();
	displaced->start();
	displaced->join();
	reference->join();

	vector<double> gradient(solution.size(),0);
	gradient[0] = (displaced->get_fitness() - reference->get_fitness())/gradient_spacing_base;
	VectorOperations::normalize_vector(gradient);
	reference->unlock();
	displaced->unlock();
	solution = line_minimization(solution,gradient);

	cout << endl << "NETWORK OPTIMIZATION STEP FINISHED" << endl << endl;
}

void SeparatedCGOptimizer::iterate_cg() {
	cout << "\n " << ++current_CG_iteration_step  << ". CG ITERATION \n \n";
		cout << solution.size() << endl;
		/*
		 * Negative tensions shall be not allowed, therefore check if a given
		 * tension vector component is negative and set it to zero.
		 */
		cout << "Modify position \n";
		for(unsigned int i=0;i<solution.size();i++)
		{
			if(solution[i]<=0.0)
				solution[i]=0.0;
		}

		/*
		 * Compute gradient (without first DOF)
		 */
		cout << "Evaluate Gradient \n";
		vector<double> g = gradient(solution);
		g[0] = 0;

		/*
		 * Don't allow positive gradients for zero tension components
		 * in order to avoid running into a negative value.
		 */
		cout << "Modify gradient \n";
		for(unsigned int i=0;i<g.size();i++)
		{
			if(solution[i]==0.0 && g[i]>0.0)
				g[i]=0.0;
		}


		vector<double> conjugated_gradient(g.size(),0);

		if(current_CG_iteration_step == 1){
			conjugated_gradient = g;
			last_conjugated_gradient = g;
			last_gradient=g;
		} else {
			double gammaZ = 0, gammaN=0, gamma=0;
			for (unsigned int i = 0; i < g.size(); i++) {
				gammaZ += (g[i]-last_gradient[i]) * g[i];
				gammaN += last_gradient[i]*last_gradient[i];
			}

			gamma = gammaZ/gammaN;

			for (unsigned int i = 0; i < g.size(); i++) {
				conjugated_gradient[i] = g[i] + gamma * last_conjugated_gradient[i];
			}
			last_gradient=g;
			last_conjugated_gradient=conjugated_gradient;
		}


		/*
		 * Apply line minimization along maximum change direction
		 */
		cout << "Apply line minimization \n";
		solution = line_minimization(solution, conjugated_gradient);
}

void SeparatedCGOptimizer::write_convergence_file(Worker * w)
{
	ofstream out ((solution_folder + "/convergence.txt").c_str(),ios::app);
	if(tot_iterations == 0) {
		initial_fitness = w->get_fitness();
		out << "Iteration <TAB> Fitness <TAB> Relative Fitness <TAB> Total Force <TAB> Total Network Force \n";
	}
	vector<double> network_solution(solution.size(),0);
	network_solution[0] = solution[0];
	Worker * w2 = worker_manager->getWorker(network_solution,-1);
	w2->work();

	out << tot_iterations << "\t";
	out << w->get_fitness() << "\t";
	out << w->get_fitness()/initial_fitness << "\t";
	out << w->get_cell_model()->get_total_force() << "\t";
	out << w2->get_cell_model()->get_total_force() << "\n";
	w2->unlock();
	out.close();
}

void SeparatedCGOptimizer::write_solution_to_file()
{
	stringstream ss;
	ss << tot_iterations;
	ofstream out ((solution_folder+"/"+ss.str()+"_solution.txt").c_str());
	for(unsigned int i=0;i<solution.size();i++)
		out << solution[i] << "\n";
	out.close();

}

void SeparatedCGOptimizer::write_FA_state_file(Worker * w)
{
	stringstream ss;
	ss << tot_iterations;
	ofstream out ((solution_folder+"/"+ss.str()+"_FA_state.txt").c_str());
	out << "# FA index <TAB> FA_area <TAB> FA_orientation <TAB> F_X [nN] <TAB> F_Y [nN] <TAB> abs(F) [nN] \n";
    unsigned int n_FAs = w->get_cell_model()->get_force_transmissions().size();
    for(unsigned int i=0;i<n_FAs;i++)
    {
    	out << i << "\t";
    	out << w->get_cell_model()->get_force_transmissions()[i]->get_shape()->get_area() << "\t";
    	out << w->get_cell_model()->get_force_transmissions()[i]->get_shape()->get_orientation_angle() << "\t";
    	out << w->get_cell_model()->get_force_transmissions()[i]->get_force()[0] << "\t";
    	out << w->get_cell_model()->get_force_transmissions()[i]->get_force()[1] << "\t";
    	out << w->get_cell_model()->get_force_transmissions()[i]->get_force_magnitude() << "\n";
    }
    out.close();

}

void SeparatedCGOptimizer::write_current_state(unsigned int step_index)
{
			/*
			 * Write intermediate solution
			 */
			stringstream ss;
			ss << step_index;

			Worker * w = worker_manager->getWorker(solution, -1);
			w->work();
			if(write_fem_output_bool) w->get_substrate_model()->write_output(solution_folder + "/" + ss.str());
			if(write_convergence_file_bool) write_convergence_file(w);
			if(write_FA_state_file_bool) write_FA_state_file(w);
			if(write_solution_bool) write_solution_to_file();
			w->unlock();
}

